import pandas as pd
from Dataset.util import Edge
from datetime import datetime


class DataSet:

    def __init__(self, end, name):
        division_radio = 0.8
        time = 0
        start_time = 0
        end_time = 0

        self.current_time = time
        self.next_time = time
        self.edge_pool = []

        if name == 'wiki':
            path = "../Dataset/wikipedia-growth.txt.gz"
            start_time = datetime.strptime('2001-02-20 06:22:33', '%Y-%m-%d %H:%M:%S')
            end_time = datetime.strptime(end, '%Y-%m-%d %H:%M:%S')
            self.df = pd.read_table(path, parse_dates=[2], header=None, infer_datetime_format=True,
                                    iterator=True, compression='gzip', error_bad_lines=False)
        elif name == 'internet':
            path = "../Dataset/internet-growth.txt.gz"
            start_time = datetime.strptime('2004-01-05', '%Y-%m-%d')
            end_time = datetime.strptime(end, '%Y-%m-%d')
            self.df = pd.read_table(path, parse_dates=[2], header=None, infer_datetime_format=True,
                                    iterator=True, compression='gzip', error_bad_lines=False)
        elif name == 'CollegeMsg':
            path = "../Dataset/CollegeMsg.txt.gz"
            start_time = datetime.fromtimestamp(1082040961)
            end_time = datetime.fromtimestamp(end)
            self.df = pd.read_table(path, sep=' ', parse_dates=[2], header=None,
                                    date_parser=lambda dates: datetime.fromtimestamp(int(dates)),
                                    iterator=True, compression='gzip', error_bad_lines=False)

        elif name == 'f2f-hypertext':
            path = "../Dataset/sociopatterns-hypertext.csv"
            start_time = datetime.fromtimestamp(1246255220)
            end_time = datetime.fromtimestamp(int(end))
            self.df = pd.read_csv(path, parse_dates=[2], header=None,
                                    date_parser=lambda dates: datetime.fromtimestamp(int(dates)),
                                    iterator=True, error_bad_lines=False)

        self.start_time = start_time
        self.end_time = end_time
        self.division_time = start_time + (end_time - start_time) * division_radio

    def read_data(self, num=50):
        trunk = self.df.get_chunk(num)
        for _, row in trunk.iterrows():
            self.edge_pool.append(Edge([row[0], row[1]], row[2]))

    def get_next(self):
        init = True
        edge_list = []
        end = False

        while True:
            pop_list = []
            if len(self.edge_pool) == 0:
                self.read_data()
            for edge in self.edge_pool:
                if edge.time != self.current_time:
                    if init:
                        self.current_time = edge.time
                        init = False
                    else:
                        self.next_time = edge.time
                        end = True
                        break
                edge_list.append(edge)
                pop_list.append(edge)
            for value in pop_list:
                self.edge_pool.remove(value)
            if end:
                break

        return edge_list
